library rbtree;

import 'package:logging/logging.dart' as logging;

// -------------------------------------------------------------------------------------------------

const DEBUG = true;
final logging.Logger logger = new logging.Logger("rbtree");

// -------------------------------------------------------------------------------------------------

const LEFT = 0;
const RIGHT = 1;

/** Returns the opposite of the specified [dir]ection. */
int opposite(int dir) {
  switch (dir) {
    case LEFT: return RIGHT;
    case RIGHT: return LEFT;
    default: throw new Error();
  }
}

class Node<K, V> {
  /** Whether this node is red (true), or black (false). */
  bool red = true;

  K key = null;
  V value = null;

  Node parent = null;
  Node left = null;
  Node right = null;

  /**
   * Virtual root above the actual tree root.
   * Actual root is the right child of super-root.
   * Actual root may be null, if the tree is empty.
   */
  Node.SuperRoot(this.right): red = false;
  Node.Red(this.parent, this.key, this.value): red = true;
  Node.Black(this.parent, this.key, this.value): red = false;

  /** Debugging representation of this tree. */
  String toString() {
    final StringBuffer sb = new StringBuffer();
    sb.write(red ? "R" : "B");
    sb.write("(key=${key}");
    if (left != null) sb.write(" l=${left}");
    if (right != null) sb.write(" r=${right}");
    sb.write(')');
    return sb.toString();
  }

  /** Reports the child of this node in the specified [dir]ection. */
  Node get(int dir) {
    switch (dir) {
      case LEFT: return left;
      case RIGHT: return right;
      default: throw new Error();
    }
  }

  /** Assigns a given [node] as a child of this node, in the specified [dir]ection. */
  void set(int dir, Node node) {
    switch (dir) {
      case LEFT: left = node; break;
      case RIGHT: right = node; break;
      default: throw new Error();
    }
  }

  /** Reports the direction (left/right) of a given immediate [child]. */
  int dirOf(Node child) {
    if (child == left) return LEFT;
    if (child == right) return RIGHT;
    throw new StateError("Node ${key} has no child ${child}");
  }

  /** Reports the sibling of this node, or null. */
  Node get sibling {
    if (parent == null) return null;
    return parent.get(opposite(parent.dirOf(this)));
  }

  /** Reports whether the specified [node] is red. */
  static bool isRed(Node node) {
    return (node != null) && node.red;
  }

  /** Reports whether the specified [node] is black. */
  static bool isBlack(Node node) {
    return (node == null) || !node.red;
  }

  void rotate(int dir) {
    switch (dir) {
      case LEFT: rotateLeft(); break;
      case RIGHT: rotateRight(); break;
      default: throw new Error();
    }
  }

  /** Rotates this node to the left. */
  void rotateLeft() {
    final p = parent;
    final pdir = p.dirOf(this);
    final a = left;
    final b = right;
    final c = b.left;
    final d = b.right;
    p.set(pdir, b);
    b.parent = p;
    b.left = this;
    this.parent = b;
    if (c != null) c.parent = this;
    this.right = c;
  }

  /** Rotates this node to the right. */
  void rotateRight() {
    final p = parent;
    final pdir = p.dirOf(this);
    final a = left;
    final b = right;
    final c = a.left;
    final d = a.right;
    p.set(pdir, a);
    a.parent = p;
    a.right = this;
    this.parent = a;
    if (d != null) d.parent = this;
    this.left = d;
  }

  void doubleRotate(int dir) {
    switch (dir) {
      case LEFT: doubleRotateLeft(); break;
      case RIGHT: doubleRotateRight(); break;
      default: throw new Error();
    }
  }

  void doubleRotateLeft() {
    right.rotateRight();
    rotateLeft();
  }

  void doubleRotateRight() {
    left.rotateLeft();
    rotateRight();
  }


  /**
   * Asserts that red-black tree properties are valid.
   * Returns the number of black nodes from this node (included) to all leaves (included).
   */
  int assertRedBlackProperties(Comparator comp, K min, K max) {
    if (red) {
      assert(!isRed(left));
      assert(!isRed(right));
    }
    if ((min != null) && (comp(min, key) >= 0)) {
      throw new StateError("${min} >= ${key}");
    }
    if ((max != null) && (comp(key, max) >= 0)) {
      throw new StateError("${key} >= ${max}");
    }
    final int blackLeft = (left == null) ? 1 : left.assertRedBlackProperties(comp, min, key);
    final int blackRight = (right == null) ? 1 : right.assertRedBlackProperties(comp, key, max);
    assert(blackLeft == blackRight);
    return red ? blackLeft : (blackLeft + 1);
  }

  /** Reports the maximum depth of this tree. */
  int depth() {
    final int ldepth = (left == null) ? 0 : left.depth();
    final int rdepth = (right == null) ? 0 : right.depth();
    return 1 + ((ldepth > rdepth) ? ldepth : rdepth);
  }

  /** Reports the minimum depth of this tree. */
  int minDepth() {
    final int ldepth = (left == null) ? 0 : left.minDepth();
    final int rdepth = (right == null) ? 0 : right.minDepth();
    return 1 + ((ldepth < rdepth) ? ldepth : rdepth);
  }

  /** Reports the size of this tree, in total number of nodes. */
  int size() {
    final int lsize = (left == null) ? 0 : left.size();
    final int rsize = (right == null) ? 0 : right.size();
    return 1 + lsize + rsize;
  }
}

// -------------------------------------------------------------------------------------------------

/** Signature of a comparator. */
typedef int Compare(Object o1, Object o2);

/** Default comparator. */
int defaultCompare(o1, o2) {
  return o1.compareTo(o2);
}

// -------------------------------------------------------------------------------------------------

/** Entry mapping one [key] to a [value]. */
class Entry<K, V> {
  final K key;
  final V value;

  /** Creates an entry with the specified [key] and the specified [value]. */
  Entry(this.key, this.value);

  /** Debugging representation of the entry. */
  String toString() {
    return "${key} -> ${value}";
  }
}

// -------------------------------------------------------------------------------------------------

/** In-order iterator of a red-black tree entries. */
class _Iterator<K, V> extends Iterator<Entry<K, V>> {
  final Node _root;

  /**
   * Stack of nodes representing the iterator's current position in the tree.
   * The top of the stack is the current entry to return.
   * null means before the first (smallest) entry.
   * empty means after the last entry.
   *
   * The visit ordering is: (node.left, node, node.right), recursive on node.left and node.right.
   */
  List<Node> _stack = null;

  _Iterator(this._root) {
  }

  /** Walks down to the deepest, left-most children of [node], pushing all of them on the stack. */
  void _stackLeftMost(Node node) {
    while (node != null) {
      _stack.add(node);
      node = node.left;
    }
  }

  bool moveNext() {
    if (_stack == null) {
      _stack = new List();
      _stackLeftMost(_root);
    } else {
      final Node node = _stack.removeLast();
      _stackLeftMost(node.right);
    }
    return !_stack.isEmpty;
  }

  Entry<K, V> get current {
    return new Entry(_stack.last.key, _stack.last.value);
  }
}

// -------------------------------------------------------------------------------------------------

class _EntryIterable<K, V> extends Iterable<Entry<K, V>> {
  final RedBlackTree<K, V> tree;
  _EntryIterable(this.tree);

  Iterator<Entry<K, V>> get iterator {
    return new _Iterator(tree._proot.right);
  }
}

/**
 * Red/black tree top-down implementation.
 * Each node maintains a pointer to its parent.
 */
class RedBlackTree<K, V> implements Map<K, V> {
  /** Comparator to order the keys. */
  final Compare _compare;

  _EntryIterable<K, V> _itEntries;

  /** Virtual grand-parent of the tree root. */
  Node<K, V> _gproot;

  /** Virtual parent of the actual root of the tree. */
  Node<K, V> _proot;

  /** Number of entries (nodes) in the tree. */
  int _nentries = 0;

  /** Creates a tree-map with a custom order. */
  factory RedBlackTree.withOrder(Compare compare) =>
      new RedBlackTree._internal(compare);

  /** Creates a new tree-map with the default order. */
  factory RedBlackTree() =>
      new RedBlackTree._internal(defaultCompare);

  /** Constructor with the default key comparator. */
  RedBlackTree._internal(this._compare) {
    _itEntries = new _EntryIterable(this);
    _proot = new Node.SuperRoot(null);
    _gproot = new Node.SuperRoot(_proot);
    _proot.parent = _gproot;
  }

  /**
   * Inserts a [value] at the specified [key] in the tree.
   * Returns the pre-existing value associated to the [key], if any, or null.
   */
  V insert(K key, V value) {
    logger.finest("insert(${key} -> ${value})");
    final V oldValue = _insert(key, value);
    logger.finest("after insertion of ${key}, tree is ${_proot.right}");
    if (DEBUG) assertRedBlackProperties();
    if (oldValue == null) _nentries += 1;
    return oldValue;
  }

  /**
   * Removes the entry with the given [key].
   * Returns the value associated to the key, or null if none existed.
   */
  V remove(K key) {
    logger.finest("remove(${key})");
    final V value = _remove(key);
    logger.finest("after removal of ${key}, tree is ${_proot.right}");
    if (DEBUG) assertRedBlackProperties();
    if (value != null) _nentries -= 1;
    return value;
  }

  /**
   * Retrieves a value by [key].
   * Returns null if the [key] is not found in the tree.
   */
  V get(K key) {
    Node<K, V> node = _proot.right;
    while (node != null) {
      final int comp = _compare(key, node.key);
      if (comp == 0) break;
      node = (comp < 0) ? node.left : node.right;
    }
    return (node != null) ? node.value : null;
  }

  V putIfAbsent(K key, V ifAbsent()) {
    final V existing = get(key);
    if (existing != null) return existing;
    final V value = ifAbsent();
    insert(key, value);
    return value;
  }

  void forEach(void function(K key, V value)) {
    for (Entry<K, V> entry in entries) {
      function(entry.key, entry.value);
    }
  }

  Iterable<K> get keys {
    return entries.map((Entry<K, V> entry) => entry.key);
  }

  Iterable<V> get values {
    return entries.map((Entry<K, V> entry) => entry.value);
  }

  Iterable<Entry<K, V>> get entries {
    return _itEntries;
  }

  bool containsKey(K key) {
    return get(key) != null;
  }

  bool containsValue(V value) {
    return values.contains(value);
  }

  V operator[](K key) {
    return get(key);
  }

  void operator []=(K key, V value) {
    insert(key, value);
  }

  void clear() {
    _proot.right = null;
    _nentries = 0;
  }

  int get length {
    return _nentries;
  }

  bool get isEmpty {
    return _nentries == 0;
  }

  // -----------------------------------------------------------------------------------------------

  V _insert(K key, V value) {
    /** Search/insertion node, guaranteed never null (except for the root case). */
    Node<K, V> inode = _proot.right;
    if (inode == null) {
      _proot.right = new Node.Black(_proot, key, value);
      return null;
    }

    final int rootComp = _compare(key, inode.key);
    if (rootComp == 0) {
      final V formerValue = inode.value;
      inode.value = value;
      return formerValue;
    }

    /** Direction from inode. */
    int dir = (rootComp < 0) ? LEFT : RIGHT;

    /** Debugging counter to keep track of the iterations. */
    int counter = 0;

    /** If an entry already exists with the specified key, this will hold the former value. */
    V formerValue = null;

    while (true) {
      logger.finest("put: iteration #${counter} with node ${inode}");
      counter += 1;

      // Loop invariant: uncle of child (ie. sibling of node) is black:
      assert(!inode.red || Node.isBlack(inode.sibling));

      final Node<K, V> child = inode.get(dir);
      if (child == null) {
        // child does not exist, node is the actual parent insertion point:
        final Node<K, V> newChild = new Node.Red(inode, key, value);
        inode.set(dir, newChild);
        if (inode.red) {
          final Node<K, V> parent = inode.parent;
          final int pdir = parent.dirOf(inode);
          if (pdir == dir) {
            logger.finest("put: single rotate parent: ${parent}");
            parent.rotate(opposite(pdir));
            parent.red = true;
            inode.red = false;
          } else {
            logger.finest("put: double rotate parent: ${parent}");
            parent.doubleRotate(opposite(pdir));
            parent.red = true;
            newChild.red = false;
          }
        }
        break;  // we are done, insert is completed.
      }  // child == null

      final int comp = _compare(key, child.key);
      if (comp == 0) {
        // child is the node with key we are looking for, override its value and return:
        formerValue = child.value;
        child.value = value;
        break;  // we are done, insert is completed
      }

      /** Direction from child. */
      final int cDir = (comp < 0) ? LEFT : RIGHT;

      // node and node.sibling cannot be both red, as a result of the previous iteration:
      if (!(Node.isRed(inode.left) && Node.isRed(inode.right))) {
        // Move down the tree and loop over:
        inode = child;
        dir = cDir;
        continue;
      }

      // Both children of node are red:
      logger.finest("put: flipping colors on node = ${inode}");
      assert(!inode.red);
      inode.red = true;
      inode.left.red = false;
      inode.right.red = false;

      final Node<K, V> parent = inode.parent;
      final int pDir = parent.dirOf(inode);

      final Node<K, V> gparent = parent.parent;
      final int gpDir = gparent.dirOf(parent);

      if (!parent.red) {
        // Move down the tree and loop over:
        inode = child;
        dir = cDir;
        continue;
      }

      // parent and node are now both red, which violates the red-black tree properties.
      if (gpDir == pDir) {
        // Single rotation of grand-parent in the opposite direction:
        logger.finest("put: single rotation on gparent ${gparent}");
        gparent.rotate(opposite(gpDir));
        parent.red = false;
        gparent.red = true;
      } else {
        // Double rotation of grand-parent in the same direction:
        logger.finest("put: double rotation on gparent ${gparent}");
        gparent.doubleRotate(opposite(gpDir));
        inode.red = false;
        gparent.red = true;
      }

      // Resume iteration directly from child, and skip the iteration: node <- child:
      final cchild = child.get(cDir);
      if (cchild == null) {
        // child does not exist, node is the actual insertion point:
        child.set(cDir, new Node.Red(child, key, value));
        if (child.red) {
          inode.rotate(opposite(dir));
          child.red = false;
          inode.red = true;
        }
        break;  // we are done, insert is completed
      }

      final ccomp = _compare(key, cchild.key);
      if (ccomp == 0) {
        // child is the node with key we are looking for, override its value and return:
        formerValue = cchild.value;
        cchild.value = value;
        break;  // insert is completed
      }
      final int ccDir = (comp < 0) ? LEFT : RIGHT;

      inode = cchild;
      dir = ccDir;

    }  // while true

    // Update/force root to black:
    _proot.right.red = false;

    return formerValue;
  }

  V _remove(K key) {
    Node<K, V> node = _proot.right;
    if (node == null) return null;

    // Flip the root node to red to initiate the removal:
    node.red = true;
    _proot.red = true;

    /** If we find the entry to be removed. */
    Node<K, V> found = null;

    /** Debugging iteration counter. */
    int counter = 0;

    while (true) {
      logger.finest("remove: iteration #${counter} on node ${node}");
      counter += 1;

      var parent = node.parent;

      // Invariant from the previous iteration:
      assert(parent.red);

      final Node<K, V> gparent = parent.parent;
      final int gpdir = gparent.dirOf(parent);

      final int pdir = parent.dirOf(node);
      final Node<K, V> sibling = parent.get(opposite(pdir));

      final int comp = _compare(key, node.key);
      if (comp == 0) {
        // node is the one we want to delete:
        found = node;
      }
      final int ndir = (comp < 0) ? LEFT : RIGHT;
      Node<K, V> child = node.get(ndir);
      final Node<K, V> csibling = node.get(opposite(ndir));

      if (Node.isRed(node.left) && Node.isRed(node.right)) {
        // node has 2 red children:
        logger.finest("remove: color flip node to red ${node}");
        node.red = true;
        node.left.red = false;
        node.right.red = false;

      } else if ((sibling != null) && Node.isBlack(node.left) && Node.isBlack(node.right)) {
        // node has 2 black children:

        final schildpdir = sibling.get(pdir);
        final schildnpdir = sibling.get(opposite(pdir));

        if (Node.isBlack(sibling.left) && Node.isBlack(sibling.right)) {
          // all four children of node and its sibling are black:
          logger.finest("remove: color flip parent to black ${parent}");
          parent.red = false;
          node.red = true;
          sibling.red = true;

        } else if (Node.isRed(schildnpdir)) {
          logger.finest("remove: rotate parent ${parent}");
          parent.rotate(pdir);

          sibling.red = true;
          parent.red = false;
          node.red = true;
          schildnpdir.red = false;

        } else if (Node.isRed(schildpdir)) {
          // sibling has at least one red child:
          logger.finest("remove: double rotate parent ${parent}");
          parent.doubleRotate(pdir);

          parent.red = false;
          node.red = true;
        }

      // Below, node has one red child:

      } else if (Node.isRed(child)) {
        // Next node to explore is red, move down the tree:
        logger.finest("remove: iterate down the tree");
        node = child;
        final comp = _compare(key, node.key);
        if (comp == 0) {
          // node is the one we want to delete:
          found = node;
        }
        final ndir = (comp < 0) ? LEFT : RIGHT;
        child = node.get(ndir);

      } else if (Node.isRed(csibling)) {
        // Next node to explore is black:
        logger.finest("remove: rotate node ${node}");
        node.rotate(ndir);
        csibling.red = false;
        node.red = true;

      } else {
        // Nothing to do, move down the tree
      }

      if (child == null) {
        assert(node.red);

        // Force tree root to black, if necessary:
        _proot.red = false;
        if (_proot.right != null) _proot.right.red = false;

        // node is the leaf to remove, assuming we found a matching entry:
        V removedValue = null;
        if (found != null) {
          removedValue = found.value;

          if (found != node) {
            logger.finest("remove: replacing ${found} with ${node}");
            found.key = node.key;
            found.value = node.value;
          }
          node.parent.set(node.parent.dirOf(node), null);
        }
        return removedValue;
      }

      // Move iterator down the tree and loop over:
      node = child;
    }  // while (true)
    throw new Error();  // dead code
  }

  /** Debugging assertion. */
  void assertRedBlackProperties() {
    if (_proot.right == null) return;
    assert(!_proot.right.red);
    _proot.right.assertRedBlackProperties(_compare, null, null);
  }
}